#include "namespace_manager.h"
#include <spdlog/spdlog.h>
#include <sys/mount.h>
#include <sys/stat.h>
#include <sys/wait.h>

#include <fcntl.h>
#include <unistd.h>
#include <sched.h>
#include <filesystem>
#include <stdexcept>
#include <signal.h>
#include <errno.h>
#include <fstream>
#include <sstream>
#include <algorithm>
#include <tuple>
#include <vector>

#ifdef __linux__
#include <sys/sysmacros.h>
// Linux特定的宏定义
#ifndef CLONE_NEWPID
#define CLONE_NEWPID 0x20000000
#endif

#ifndef CLONE_NEWNS
#define CLONE_NEWNS 0x00020000
#endif

#ifndef CLONE_NEWNET
#define CLONE_NEWNET 0x40000000
#endif

#ifndef CLONE_NEWUTS
#define CLONE_NEWUTS 0x04000000
#endif

#ifndef CLONE_NEWIPC
#define CLONE_NEWIPC 0x08000000
#endif

#ifndef MS_REC
#define MS_REC 0x4000
#endif

#ifndef MS_PRIVATE
#define MS_PRIVATE (1<<18)
#endif

#ifndef MS_BIND
#define MS_BIND 0x1000
#endif

#ifndef MS_REMOUNT
#define MS_REMOUNT 0x20
#endif

#ifndef MS_RDONLY
#define MS_RDONLY 0x1
#endif

#ifndef MNT_DETACH
#define MNT_DETACH 0x2
#endif
#endif

namespace judge {

namespace fs = std::filesystem;

NamespaceManager::NamespaceManager() : root_dir_("/tmp/judge/ns_root") {
    if (root_dir_.empty()) {
        throw std::runtime_error("Root directory cannot be empty");
    }
}

NamespaceManager::NamespaceManager(const std::string& root_dir) : root_dir_(root_dir) {
    if (root_dir_.empty()) {
        throw std::runtime_error("Root directory cannot be empty");
    }
}

NamespaceManager::~NamespaceManager() {
    cleanup();
}

bool NamespaceManager::initializeNamespaces() {
#ifdef __linux__
    try {
        // 创建所有命名空间
        unsigned long ns_flags = CLONE_NEWNS | CLONE_NEWPID | CLONE_NEWUTS | CLONE_NEWIPC | CLONE_NEWNET;
        if (unshare(ns_flags) == -1) {
            throw std::runtime_error("Failed to unshare namespaces: " + std::string(strerror(errno)));
        }

        // 确保mount命名空间是私有的
        if (mount(nullptr, "/", nullptr, MS_REC | MS_PRIVATE, nullptr) == -1) {
            throw std::runtime_error("Failed to make root mount private: " + std::string(strerror(errno)));
        }

        return true;
    } catch (const std::exception& e) {
        spdlog::error("Failed to initialize namespaces: {}", e.what());
        return false;
    }
#else
    return true;
#endif
}

void NamespaceManager::setRootDirectory(const std::string& path) {
    if (path.empty()) {
        throw std::invalid_argument("Root directory cannot be empty");
    }
    root_dir_ = path;
}

void NamespaceManager::addMountPoint(const std::string& source, const std::string& target,
                                     const std::string& fstype, unsigned long flags) {
    mount_points_.push_back({source, target, fstype, flags});
}

// Extend initializeFilesystem to mount user-specified mount points after basic setup
void NamespaceManager::initializeFilesystem() {
#ifdef __linux__
    try {
        // 确保根目录存在
        if (!fs::exists(root_dir_)) {
            fs::create_directories(root_dir_);
        }

        // 创建基本目录结构
        std::vector<std::string> dirs = {
            "/bin", "/lib", "/lib64", "/usr", "/proc", "/dev", "/sys", "/etc", "/tmp"
        };
        for (const auto& dir : dirs) {
            fs::create_directories(root_dir_ + dir);
        }

        // 挂载特殊文件系统
        const std::vector<std::tuple<std::string, std::string, std::string, unsigned long>> mounts = {
            {"/proc", root_dir_ + "/proc", "proc", MS_NOSUID | MS_NODEV | MS_NOEXEC},
            {"/dev", root_dir_ + "/dev", "tmpfs", MS_NOSUID},
            {"/sys", root_dir_ + "/sys", "sysfs", MS_NOSUID | MS_NODEV | MS_NOEXEC | MS_RDONLY}
        };

        for (const auto& [src, dst, type, flags] : mounts) {
            if (mount(type.c_str(), dst.c_str(), type.c_str(), flags, nullptr) == -1) {
                spdlog::warn("Failed to mount {} at {}: {}", type, dst, strerror(errno));
            }
        }

        // 创建必要的设备节点
        const std::vector<std::tuple<std::string, mode_t, dev_t>> devices = {
            {"/dev/null", S_IFCHR | 0666, makedev(1, 3)},
            {"/dev/zero", S_IFCHR | 0666, makedev(1, 5)},
            {"/dev/random", S_IFCHR | 0666, makedev(1, 8)},
            {"/dev/urandom", S_IFCHR | 0666, makedev(1, 9)}
        };

        for (const auto& [path, mode, dev] : devices) {
            std::string full_path = root_dir_ + path;
            if (mknod(full_path.c_str(), mode, dev) == -1 && errno != EEXIST) {
                spdlog::warn("Failed to create device node {}: {}", full_path, strerror(errno));
            }
            if (chmod(full_path.c_str(), mode & 0777) == -1) {
                spdlog::warn("Failed to set permissions for {}: {}", full_path, strerror(errno));
            }
        }

        // Mount user-defined mount points (read-only or bind as requested)
        for (const auto& mp : mount_points_) {
            std::string dst = root_dir_ + mp.target;
            if (!fs::exists(dst)) {
                fs::create_directories(dst);
            }
            const char* type_ptr = mp.fstype.empty() ? nullptr : mp.fstype.c_str();
            if (mount(mp.source.c_str(), dst.c_str(), type_ptr, mp.flags, nullptr) == -1) {
                spdlog::warn("Failed to mount {} to {}: {}", mp.source, dst, strerror(errno));
            }
        }

        // 切换根目录
        if (chroot(root_dir_.c_str()) == -1) {
            throw std::runtime_error("Failed to change root: " + std::string(strerror(errno)));
        }
        if (chdir("/") == -1) {
            throw std::runtime_error("Failed to change directory to /: " + std::string(strerror(errno)));
        }

    } catch (const std::exception& e) {
        spdlog::error("Failed to initialize filesystem: {}", e.what());
        throw;
    }
#endif
}

void NamespaceManager::initialize() {
    if (!initializeNamespaces()) {
        throw std::runtime_error("Failed to initialize namespaces");
    }
    initializeFilesystem();
}

void NamespaceManager::cleanup() {
#ifdef __linux__
    try {
        // 卸载所有挂载点
        std::vector<std::string> mounts = {"/proc", "/dev", "/sys"};
        for (const auto& mount : mounts) {
            std::string full_path = root_dir_ + mount;
            if (fs::exists(full_path)) {
                if (umount2(full_path.c_str(), MNT_DETACH) == -1) {
                    spdlog::warn("Failed to unmount {}: {}", full_path, strerror(errno));
                }
            }
        }

        // 删除根目录
        if (!root_dir_.empty() && fs::exists(root_dir_)) {
            fs::remove_all(root_dir_);
        }
    } catch (const std::exception& e) {
        spdlog::error("Failed to cleanup namespace: {}", e.what());
    }
#endif
}

} // namespace judge